from __future__ import with_statement
"""
Created on 26-Feb-2010

@author: dexen
"""

import time
import threading
from types import DictType

import rpyc

from dexen.system.common.base import Action
from dexen.system.common.meta import JobMeta
from dexen.system.common.utils import NameMgr
from dexen.system.common.service import DexenAddr
from dexen.system.server.action_mgr import ActionMgr
from dexen.system.server.proxy import NodeProxy, MasterProxy
from dexen.system.common.excepts import AllNodeBusyException
from dexen.system.common.transact import SlaveTransact, MasterRecvTransact
from dexen.system.common.states import (SlaveState, JobState, MasterState,
                                        M_FINISHED, M_STOPPED)

#===============================================================================
# Constants
#===============================================================================
INF = 100000   

#===============================================================================
# ServerCore class
#===============================================================================
class ServerCore(threading.Thread):
    def __init__(self, reg_addr):
        threading.Thread.__init__(self)
        self.setDaemon(True)
        self.name = "ServerCore"
        
        self.action_mgrs = {} # job_id:ActionMgr    
        self.nodes = {} # node_name:NodeProxy
        self.slaves = {} # slave_name:SlaveState
        self.jobs = {} # job_id:JobMeta
        
        self.master_namer = NameMgr("MASTER")    
        self.slave_namer = NameMgr("SLAVE")            
        self.cur_job_id = 1 # current available id
        self.reg_addr = reg_addr # The DexenAddr(ip, port) address of the registry
        self.lock = threading.RLock()

    def _recover_slaves(self, master):
        assert isinstance(master, MasterProxy)
        for slave_name in master.get_scheduled_slaves():
            slave = self.slaves[slave_name]
            assert isinstance(slave, SlaveState)
            slave.idle()
        master.reset_scheduled_slaves()

    def _update_masters(self):
        for node in self.nodes.values():
            assert isinstance(node, NodeProxy)
            if node.is_crashed():
                #print "node:", node.name, " crashed." 
                continue
            node.ping()
            for job_id, master in node.masters.items():
                assert isinstance(master, MasterProxy)
                
                if not master.is_running(): continue
                
                master.update()
                actions = master.get_actions()
                act_mgr = self.action_mgrs[job_id]
                assert isinstance(act_mgr, ActionMgr)
                act_mgr.add_actions(actions)
                
                if master.is_crashed(): 
                    self._recover_slaves(master)
                    continue
                #TODO: Handle the case when master is running and node is crashed.
                idle_slave = None
                for action in actions:
                    assert isinstance(action, Action)
                    transact = action.get_transact()               
                    if action.is_slave():
                        assert isinstance(transact, SlaveTransact)
                        slave = self.slaves[transact.slave_name]
                        assert isinstance(slave, SlaveState)
                        slave.update(transact)
                        if slave.is_idle():
                            idle_slave = slave
                    if action.is_master():
                        assert isinstance(transact, MasterRecvTransact)
                        master_state = action.get_master_state()
                        assert isinstance(master_state, MasterState)
                        if master_state.is_finished():
                            node.stop_master(job_id, M_FINISHED)
                        
                if master.is_running() and master.can_schedule():
                    if idle_slave == None:
                        idle_slave = node.get_free_slave()
                        if idle_slave == None: 
                            for node in self.nodes.values():
                                assert isinstance(node, NodeProxy)
                                idle_slave = node.get_free_slave()
                                if idle_slave: break
                    if idle_slave != None:
                        assert isinstance(idle_slave, SlaveState)
                        idle_slave.busy()
                        master.schedule(idle_slave.name)                  
                            
    def _create_slaves(self, node):
        assert isinstance(node, NodeProxy)
        n_slots = node.n_slave_slots()
        for _ in range(n_slots):
            name = self.slave_namer.get_name()
            node.create_slave(name)
            if node.is_crashed(): break
            self.slaves[name] = node.get_slave(name)
            
    def _discover_nodes(self):
        """ 
            The new_nodes and dead_nodes are of type dict.
            key: the name of the nodes
            value: the timestamp of the creation of node
        """
        new_nodes = self.reg_conn.root.get_new_nodes()
        new_nodes = rpyc.classic.obtain(new_nodes)
        assert isinstance(new_nodes, DictType)
        for name, addr in new_nodes.items():
            self.nodes[name] = NodeProxy(name, addr, self.reg_addr) 
            self._create_slaves(self.nodes[name])

    def _get_free_node(self):
        """
        For creating master
        TODO: Handle the case when all the slots are taken
        """
        for node in self.nodes.values():
            assert isinstance(node, NodeProxy)
            if node.is_running() and node.has_master_slot():
                return node              
        raise AllNodeBusyException("There is no available Node.") 
                            
    def run(self):
        """
            This is the starting point of core thread running in parallel 
            with service thread that serves clients to submit jobs,
            etc..
        """
        reg_addr = self.reg_addr
        assert isinstance(reg_addr, DexenAddr)
        self.reg_conn = reg_addr.get_conn()
        while True:
            try:
                with self.lock:
                    self._discover_nodes()
                    self._update_masters()
                time.sleep(0.1)
            except KeyboardInterrupt:
                print "Keyboard interrupted."
                break
            except EOFError:
                print "EOFError"
                break
    
    #===========================================================================
    # Methods for service
    #===========================================================================
    def get_actions(self, job_id, act_id):
        with self.lock:
            act_mgr = self.action_mgrs[job_id]
            assert isinstance(act_mgr, ActionMgr)
            return act_mgr.get_actions(act_id)
    
    def get_client_config_file(self, job_id):
        with self.lock:
            job_meta = self.jobs[job_id]
            assert isinstance(job_meta, JobMeta)
            return job_meta.get_config_file()
    
    def get_job_states(self):
        job_states = []
        with self.lock:
            for job_id, job_meta in self.jobs.items():
                assert isinstance(job_meta, JobMeta)
                master = self._get_master_proxy(job_id)
                assert isinstance(master, MasterProxy)
                master_state = master.get_master_state()
                js = JobState(job_meta, master_state)
                job_states.append(js)
            return job_states
    
    def get_job_state(self, job_id):
        master_proxy = self._get_master_proxy(job_id)
        assert isinstance(master_proxy, MasterProxy)
        master_state = master_proxy.get_master_state()
        job_meta = self.jobs[job_id]
        return JobState(job_meta, master_state)
    
    def start_master(self, schema_zip, job_name):
        with self.lock:
            master_name = self.master_namer.get_name()
            while True:
                node = self._get_free_node()
                assert isinstance(node, NodeProxy)
                job_meta = JobMeta(self.cur_job_id, job_name, node.name, 
                                   schema_zip, time.time())
                node.create_master(master_name, job_meta) # creates the master proxy
                if not node.is_crashed(): # means creation of master is successful
                    node.start_master(job_meta.get_id()) # start running the job
                    break 
            self.jobs[job_meta.get_id()] = job_meta
            self.action_mgrs[job_meta.get_id()] = ActionMgr(job_meta)
            self.cur_job_id += 1            
 
    def _get_node(self, job_id):
        job_meta = self.jobs[job_id]
        assert isinstance(job_meta, JobMeta)
        return self.nodes[job_meta.node_name]
    
    def _get_master_proxy(self, job_id):
        node = self._get_node(job_id)
        assert isinstance(node, NodeProxy)
        return node.get_master(job_id)
    
    def stop_master(self, job_id): # called by the user (cli)
        with self.lock:
            node = self._get_node(job_id)
            assert isinstance(node, NodeProxy)
            node.stop_master(job_id, M_STOPPED)
            master = node.get_master(job_id)
            self._recover_slaves(master)

    def del_master(self, job_id):
        with self.lock:
            node = self._get_node(job_id)
            assert isinstance(node, NodeProxy)
            node.del_master(job_id) # if master is running, fire an exception
            action_mgr = self.action_mgrs[job_id]
            assert isinstance(action_mgr, ActionMgr)
            action_mgr.on_delete()
            del self.action_mgrs[job_id]
            del self.jobs[job_id]
            
    def pause_master(self, job_id):
        pass
        
    def __repr__(self):
        t = ""
        for name, node in self.nodes.items():
            t += "%s:\n%r\n"%(name, node)
        return t
